"""
Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
SPDX-License-Identifier: Apache-2.0
"""

import logging

from graph_notebook.network.EventfulNetwork import EventfulNetwork, DEFAULT_GRP, DEPTH_GRP_KEY, DEFAULT_RAW_GRP_KEY
from networkx import MultiDiGraph

logging.basicConfig()
logger = logging.getLogger(__file__)

DEFAULT_LABEL_MAX_LENGTH = 10
ENTITY_KEY = "~entityType"
ID_KEY = "~id"
START_KEY = "~start"
END_KEY = "~end"
PROPERTIES_KEY = "~properties"
VERTEX_TYPE_KEY = "~entityType"
EDGE_TYPE_KEY = "~type"
LABEL_KEY = "~labels"
NODE_ENTITY_TYPE = 'node'
REL_ENTITY_TYPE = 'relationship'


class OCNetwork(EventfulNetwork):
    """OCNetwork extends the EventfulNetwork class and uses the add_results method to parse any response that returns
    nodes/relationships as part (or all) of the response.
    """

    def __init__(self, graph: MultiDiGraph = None, callbacks=None, label_max_length=DEFAULT_LABEL_MAX_LENGTH,
                 edge_label_max_length=DEFAULT_LABEL_MAX_LENGTH, group_by_property=LABEL_KEY,
                 display_property=LABEL_KEY, edge_display_property=EDGE_TYPE_KEY,
                 tooltip_property=None, edge_tooltip_property=None,
                 ignore_groups=False, 
                 group_by_depth=False, group_by_raw=False):
        if graph is None:
            graph = MultiDiGraph()
        if group_by_depth:
            group_by_property = DEPTH_GRP_KEY
        super().__init__(graph, callbacks, label_max_length, edge_label_max_length, group_by_property,
                         display_property, edge_display_property, tooltip_property, edge_tooltip_property,
                         ignore_groups, group_by_raw)

    def get_node_property_value(self, node: dict, props: dict, title, custom_property):
        try:
            if isinstance(custom_property, dict):
                if isinstance(custom_property[title], tuple):
                    if custom_property[title][0] in props and \
                            isinstance(props[custom_property[title][0]], list) and \
                            len(props[custom_property[title][0]]) >= 2:
                        label = props[custom_property[title][0]][custom_property[title][1]]
                    else:
                        label = title
                elif custom_property[title] in props:
                    label = props[custom_property[title]]
                elif LABEL_KEY in props:
                    label = props[LABEL_KEY]
                else:
                    label = props
            elif isinstance(custom_property, tuple) and custom_property[0] in props:
                if isinstance(props[custom_property[0]], list) and len(props[custom_property[0]]) >= 2:
                    label = props[custom_property[0]][custom_property[1]]
                else:
                    label = title
            elif custom_property == ID_KEY:
                label = node[ID_KEY]
            elif custom_property == LABEL_KEY:
                label = node[LABEL_KEY]
            elif custom_property == VERTEX_TYPE_KEY:
                label = node[VERTEX_TYPE_KEY]
            elif custom_property in props:
                label = props[custom_property]
            else:
                label = title
        except (KeyError, IndexError) as e:
            logger.debug(e)
            label = title

        return label

    def get_edge_property_value(self, data: dict, rel: dict, custom_property):
        if custom_property is not EDGE_TYPE_KEY:
            try:
                if isinstance(custom_property, dict):
                    if isinstance(custom_property[data['label']], tuple) and \
                            custom_property[data['label']][0] in data['properties']:
                        if isinstance(data['properties'][custom_property[data['label']][0]], list) and \
                                len(data['properties'][custom_property[data['label']][0]]) >= 2:
                            display_label = str(data['properties'][custom_property[data['label']][0]]
                                                [custom_property[data['label']][1]])
                        else:
                            display_label = rel[EDGE_TYPE_KEY]
                    else:
                        display_label = data['properties'][custom_property[rel[EDGE_TYPE_KEY]]]
                elif isinstance(custom_property, tuple) and \
                        custom_property[0] in data['properties']:
                    display_label = str(data['properties'][custom_property[0]]
                                        [custom_property[1]])
                else:
                    display_label = data['properties'][custom_property]
            except (KeyError, IndexError, TypeError) as e:
                logger.debug(e)
                display_label = rel[EDGE_TYPE_KEY]
        else:
            display_label = rel[EDGE_TYPE_KEY]

        return display_label

    def parse_node(self, node: dict, path_index: int = -2):
        """This parses the node parameter and adds the node to the network diagram

        Args:
            node (dict): The node dictionary to parse
            path_index: Position of the element in the path traversal order
        """

        depth_group = "__DEPTH-" + str(path_index//2) + "__"

        # generate placeholder tooltip from label; if not present, amalgamate node property values instead
        if LABEL_KEY in node.keys():
            if len(node[LABEL_KEY]) > 0:
                title_plc = node[LABEL_KEY][0]
                create_title_placeholder = False
            else:
                create_title_placeholder = True
        else:
            create_title_placeholder = True

        if create_title_placeholder:
            title_plc = ""
            for key in node:
                title_plc += str(node[key])

        if not isinstance(self.group_by_property, dict):  # Handle string format group_by
            try:
                if self.group_by_property == DEPTH_GRP_KEY:
                    group = depth_group
                elif self.group_by_property == DEFAULT_RAW_GRP_KEY:
                    group = str(node)
                elif self.group_by_property == LABEL_KEY and len(node[LABEL_KEY]) > 0:
                    group = node[LABEL_KEY][0]
                elif self.group_by_property == ID_KEY:
                    group = node[ID_KEY]
                elif self.group_by_property in node[PROPERTIES_KEY]:
                    group = node[PROPERTIES_KEY][self.group_by_property]
                else:
                    group = DEFAULT_GRP
            except KeyError:
                group = DEFAULT_GRP
        else:  # handle dict format group_by
            try:
                if str(node[LABEL_KEY][0]) in self.group_by_property and len(node[LABEL_KEY]) > 0:
                    key = node[LABEL_KEY][0]
                    if self.group_by_property[key] == DEPTH_GRP_KEY:
                        group = depth_group
                    elif self.group_by_property[key] == DEFAULT_RAW_GRP_KEY:
                        group = str(node)
                    elif self.group_by_property[key] == LABEL_KEY:
                        group = node[LABEL_KEY][0]
                    elif self.group_by_property[key] == ID_KEY:
                        group = node[ID_KEY]
                    else:
                        group = node[PROPERTIES_KEY][self.group_by_property[key]]
                else:
                    group = DEFAULT_GRP
            except KeyError:
                group = DEFAULT_GRP

        props = self.flatten(node)
        label = self.get_node_property_value(node, props, title_plc, self.display_property)
        if not label:
            label = node[ID_KEY]
        title, label = self.strip_and_truncate_label_and_title(label, self.label_max_length)
        if self.tooltip_property and self.tooltip_property != self.display_property:
            title, label_plc = self.strip_and_truncate_label_and_title(
                self.get_node_property_value(node, props, title_plc, self.tooltip_property))
        data = {'properties': props, 'label': label, 'title': title, 'group': group}
        if self.ignore_groups:
            data['group'] = DEFAULT_GRP
        self.add_node(node[ID_KEY], data)
    
    def parse_rel(self, rel):
        data = {'properties': self.flatten(rel), 'label': rel[EDGE_TYPE_KEY], 'title': rel[EDGE_TYPE_KEY]}
        display_label = self.get_edge_property_value(data, rel, self.edge_display_property)
        edge_title, edge_label = self.strip_and_truncate_label_and_title(display_label, self.edge_label_max_length)
        if self.edge_tooltip_property and self.edge_tooltip_property != self.edge_display_property:
            edge_title = self.get_edge_property_value(data, rel, self.edge_tooltip_property)
        data['title'] = edge_title
        data['label'] = edge_label
        self.add_edge(from_id=rel[START_KEY], to_id=rel[END_KEY], edge_id=rel[ID_KEY], label=edge_label,
                      title=edge_title, data=data)

    def process_result(self, res: dict, path_index: int = -2):
        """Determines the type of element passed in and processes it appropriately

        Args:
            res (dict): The dictionary to parse
            path_index: Position of the element in the path traversal order
        """
        if ENTITY_KEY in res:
            if res[ENTITY_KEY] == NODE_ENTITY_TYPE:
                self.parse_node(res, path_index)
            else:
                self.parse_rel(res)

    def add_results(self, results):
        """Adds the results parameter to the network

        Args:
            results (Object): Determines the type of the object and processes it appropriately
        """
        for res in results["results"]:
            if type(res) is dict:
                for k in res.keys():
                    if type(res[k]) is dict:
                        self.process_result(res[k]) 
                    elif type(res[k]) is list:
                        for path_index, res_sublist in enumerate(res[k]):
                            try:
                                self.process_result(res_sublist, path_index)
                            except TypeError as e:
                                logger.debug(f'Property {res_sublist} in list results set is invalid, skipping')
                                logger.debug(f'Error: {e}')
                                continue
